import gym
from gym import spaces
import numpy as np


class SimplePVZEnv(gym.Env):
    def __init__(self):
        # 定义状态空间：假设阳光数量、植物种类和位置（简化为0 - 9的位置）、僵尸位置（同样简化）
        self.observation_space = spaces.Dict({
           'sunlight': spaces.Discrete(1000),
            'plant_locations': spaces.MultiDiscrete([10] * 5),
            'zombie_location': spaces.Discrete(10)
        })
        # 定义动作空间：选择植物（0 - 3 种植物）和放置位置（0 - 9）
        self.action_space = spaces.Tuple((
            spaces.Discrete(4),
            spaces.Discrete(10)
        ))
        self.reset()

    def reset(self):
        self.sunlight = 500
        self.plant_locations = np.array([-1] * 5)
        self.zombie_location = 0
        return {
           'sunlight': self.sunlight,
            'plant_locations': self.plant_locations,
            'zombie_location': self.zombie_location
        }

    def step(self, action):
        plant_type, plant_location = action
        reward = 0
        done = False

        # 检查阳光是否足够放置植物
        if plant_type == 0 and self.sunlight >= 50:
            self.sunlight -= 50
            for i in range(len(self.plant_locations)):
                if self.plant_locations[i] == -1:
                    self.plant_locations[i] = plant_location
                    break

        # 僵尸移动
        self.zombie_location += 1
        if self.zombie_location >= 10:
            reward = -100
            done = True
        else:
            for loc in self.plant_locations:
                if loc == self.zombie_location:
                    reward = 50
                    self.plant_locations = np.array([-1] * 5)

        return {
           'sunlight': self.sunlight,
            'plant_locations': self.plant_locations,
            'zombie_location': self.zombie_location
        }, reward, done, {}


# 以下是简单的随机策略测试
env = SimplePVZEnv()
observation = env.reset()
total_reward = 0
for _ in range(100):
    action = env.action_space.sample()
    observation, reward, done, _ = env.step(action)
    total_reward += reward
    if done:
        break
print(f"Total reward: {total_reward}")